import rospy
from sensor_msgs.msg import NavSatFix, Imu
from nav_msgs.msg import Odometry
from tf.transformations import euler_from_quaternion
from std_msgs.msg import Float32
from gazebo_msgs.srv import SetModelState
from gazebo_msgs.msg import ModelState
import math
import threading
from autopilot import navigate_to_target
from enu_xyz import GNC_COORDINATE
from queue import Queue
import numpy as np

import torch
import time
from DDPG_ import DDPG
from DDPG_ import PolicyNetwork
import matplotlib.pyplot as plt
# 在全局范围内创建一个队列
boat_pos_queue = Queue()
boat_yaw_queue = Queue()
boat_vel_queue = Queue()
boat_speed_queue = Queue()
M_PI = math.pi/3
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")

# 输出船的舵角和推力
class propulsion_VRX:
    def __init__(self, propeller_mid, servo_mid):
        self.last_propeller = 0
        self.last_servo = 0
        self.propeller_mid = propeller_mid
        self.servo_mid = servo_mid
        self.propeller_k = (1 - (-1)) / (100.0 - (-100.0))
        self.servo_k = M_PI / (100.0 - (-100.0))

        self.left_angle = rospy.Publisher('/wamv/thrusters/left_thrust_angle', Float32, queue_size=1)
        self.left_cmd = rospy.Publisher('/wamv/thrusters/left_thrust_cmd', Float32, queue_size=1)
        self.right_angle = rospy.Publisher('/wamv/thrusters/right_thrust_angle', Float32, queue_size=1)
        self.right_cmd = rospy.Publisher('/wamv/thrusters/right_thrust_cmd', Float32, queue_size=1)

    def get_trans_propeller(self, propeller):
        propeller_trans = propeller
        # if propeller >= 100:
        #     propeller_trans = 1
        # elif propeller <= -100:
        #     propeller_trans = -1
        # elif propeller == 0:
        #     propeller_trans = self.propeller_mid
        # else:
        #     propeller_trans = self.propeller_k * (propeller + self.propeller_mid)

        if propeller_trans > 1:
            propeller_trans = 1
        elif propeller_trans < -1:
            propeller_trans = -1

        return propeller_trans

    def get_trans_servo(self, servo):
        servo = 100*servo
        if servo >= 100:
            servo_trans = M_PI / 2
        elif servo <= -100:
            servo_trans = -M_PI / 2
        else:
            servo_trans = self.servo_k * (servo + self.servo_mid)

        if servo_trans > M_PI / 2:
            servo_trans = M_PI / 2
        elif servo_trans < -M_PI / 2:
            servo_trans = -M_PI / 2

        return servo_trans

    def manual_act(self, propeller_L, propeller_R, servo):
        power_L = self.get_trans_propeller(propeller_L)
        power_R = self.get_trans_propeller(propeller_R)
        turn = self.get_trans_servo(servo)

        # self.last_propeller = propeller
        # self.last_servo = servo

        f32_value = Float32()
        f32_value_ = Float32()
        f32_value.data = power_L
        f32_value_.data = power_R
        self.left_cmd.publish(f32_value)
        self.right_cmd.publish(f32_value_)

        f32_value.data = turn
        self.left_angle.publish(f32_value)
        self.right_angle.publish(f32_value)

# vrx接收惯导数据
class ROSVRXInfo:
    def __init__(self, update_func_gps, update_func_ea, update_func_ground_speed,
                 update_func_gyro, update_func_acc, update_handle):
        self.m_update_func_gps = update_func_gps
        self.m_update_func_ea = update_func_ea
        self.m_update_func_ground_speed = update_func_ground_speed
        self.m_update_func_gyro = update_func_gyro
        self.m_update_func_acc = update_func_acc
        self.m_update_handle = update_handle
        self.m_quit = False

    def NavSatFix_cb(self, fix):
        if self.m_update_func_gps is not None:
            self.m_update_func_gps(self.m_update_handle, fix.latitude, fix.longitude, fix.altitude)

    def Odometry_cb(self, odom):
        if self.m_update_func_ea is not None:
            quat = (
                odom.pose.pose.orientation.x,
                odom.pose.pose.orientation.y,
                odom.pose.pose.orientation.z,
                odom.pose.pose.orientation.w
            )
            (roll, pitch, yaw) = euler_from_quaternion(quat)
            yaw = (yaw + 2*math.pi) % (2*math.pi) # Ensure yaw is in [0, 2*pi)
            yaw = math.degrees(yaw)
            self.m_update_func_ea(self.m_update_handle, yaw, math.degrees(pitch), math.degrees(roll))

        if self.m_update_func_ground_speed is not None:
            boat_v = odom.twist.twist  # x是船头方向
            ground_speed = math.sqrt(boat_v.linear.x**2 + boat_v.linear.y**2)
            self.m_update_func_ground_speed(self.m_update_handle, ground_speed , boat_v.linear.x , boat_v.linear.y)

        if self.m_update_func_gyro is not None:
            self.m_update_func_gyro(self.m_update_handle,
                                    math.degrees(odom.twist.twist.angular.x),
                                    math.degrees(odom.twist.twist.angular.y),
                                    math.degrees(odom.twist.twist.angular.z))

    def IMU_cb(self, imu):
        if self.m_update_func_acc is not None:
            self.m_update_func_acc(self.m_update_handle, imu.linear_acceleration.x, imu.linear_acceleration.y,
                                   imu.linear_acceleration.z)

       # rospy.loginfo("Latitude: {:.6f}, Longitude: {:.6f}, Altitude: {:.6f}".format(imu.linear_acceleration.z,imu.linear_acceleration.y, imu.linear_acceleration.z))
    def proc_ROSVRX_info(self):

        rospy.loginfo("ROSVRX_info node started")
        # start_time = rospy.Time.now()
        # def get_current_time():
        #     return (rospy.Time.now() - start_time).to_sec()

        rospy.Subscriber('/wamv/robot_localization/gps/filtered', NavSatFix, self.NavSatFix_cb)
        rospy.Subscriber('/wamv/robot_localization/odometry/filtered', Odometry, self.Odometry_cb)
        rospy.Subscriber('/wamv/sensors/imu/imu/data', Imu, self.IMU_cb)

        # while not self.m_quit and not rospy.is_shutdown():
        #     current_time = get_current_time()
        #     rospy.loginfo("Received data at {:.6f} seconds".format(current_time))

    def ROSVRX_info_create(self):
        rospy.loginfo("ROSVRX_info create done! handle ptr {}.".format(id(self)))
        self.proc_ROSVRX_info()

def plot(frame_idx, rewards):
    plt.figure(figsize=(5, 5))
    # plt.subplot(131)
    plt.title('frame %s. reward: %s' % (frame_idx, rewards[-1]))
    plt.plot(rewards)
    plt.show()


def calculate_metrics(yaw, forward_v, lateral_v, origin_x, origin_y, hangdian2_x, hangdian2_y, boat_x, boat_y):
    yaw_rad = math.radians(yaw)

    # 计算船在指定向量方向上的速度投影
    # vector_hangdian = np.array([hangdian2_x - origin_x, hangdian2_y - origin_y])

    # v_x = np.array([forward_v * math.cos(yaw_rad), forward_v * math.sin(yaw_rad)])
    # v_y = np.array([lateral_v * math.cos(yaw_rad), lateral_v * math.sin(yaw_rad)])
    # velocity_projection = (np.dot(v_x, vector_hangdian) / math.sqrt(np.dot(vector_hangdian, vector_hangdian)) )
    #                        # + np.dot(v_y, vector_hangdian) / math.sqrt(np.dot(vector_hangdian, vector_hangdian)))

    # 计算船到指定直线的距离
    line_vector_x = hangdian2_x - origin_x
    line_vector_y = hangdian2_y - origin_y
    boat_vector_x = boat_x - origin_x
    boat_vector_y = boat_y - origin_y
    distance_to_line = (line_vector_x * boat_vector_y - line_vector_y * boat_vector_x) / math.sqrt(line_vector_x**2 + line_vector_y**2)
    # 计算矢量场制导律的角度
    k = 0.3
    ang_ = -math.atan(k * distance_to_line) + math.atan2(line_vector_y, line_vector_x)

    # 计算ang_和yaw_rad的夹角，限制在-pi到pi之间
    angle_error = ang_ - yaw_rad
    # print('slc:',math.degrees(ang_),'yaw:',math.degrees(yaw_rad))
    # 将角度误差调整到 -π 到 π 之间
    angle_error = (angle_error + math.pi) % (2 * math.pi) - math.pi

    # 计算船的艏向角度与指定向量的角度误差
    # vector_angle = math.radians(math.atan2(vector_y, vector_x))
    # if vector_angle < 0 :
    #     vector_angle += 2 * math.pi
    # angle_error = vector_angle - yaw_rad
    # if angle_error > math.pi:
    #     angle_error = angle_error - 2 * math.pi
    # elif angle_error < -math.pi:
    #     angle_error = angle_error + 2 * math.pi
    # angle_error = angle_error

    return distance_to_line, angle_error

def random_target():
    target_x = 0
    target_y = 0
    while (target_x ** 2 + target_y ** 2) < 36:
        target_x = np.random.uniform(-10, 10)
        target_y = np.random.uniform(-10, 10)
    return target_x, target_y

def goto_point(given_point):
    # 等待服务可用
    rospy.wait_for_service('/gazebo/set_model_state')
    try:
        # 创建服务代理
        set_state = rospy.ServiceProxy('/gazebo/set_model_state', SetModelState)

        # 构造状态消息
        state_msg = ModelState()
        state_msg.model_name = 'wamv'
        state_msg.pose.position.x = given_point[0]
        state_msg.pose.position.y = given_point[1]
        state_msg.pose.position.z = 0.0  # 假设 z 坐标为 0
        # 保持默认朝向（或者根据需要设置四元数）
        state_msg.pose.orientation.x = 0.0
        state_msg.pose.orientation.y = 0.0
        state_msg.pose.orientation.z = math.sin(math.pi / 4)  # 0.707
        state_msg.pose.orientation.w = math.cos(math.pi / 4)  # 0.707
        state_msg.reference_frame = 'world'

        # 调用服务进行状态设置（瞬移）
        resp = set_state(state_msg)
        print("无人船已瞬移到目标点：", given_point)
    except rospy.ServiceException as e:
        print("调用 /gazebo/set_model_state 服务失败: %s" % e)

def plot_boat_coordinates(boat_coordinates, string):
    # 提取 x 和 y 坐标
    x_coordinates, y_coordinates = zip(*boat_coordinates)

    # 绘制船的轨迹
    plt.plot(x_coordinates, y_coordinates, marker='o', linestyle='-')
    plt.axis('auto')  # 自动调整坐标轴范围
    # 设置图形属性
    plt.title(string)
    plt.xlabel('X Coordinate')
    plt.ylabel('Y Coordinate')
    plt.grid(True)

    # 显示图形
    plt.show()

if __name__ == "__main__":

    gnc_coordinate = GNC_COORDINATE()

    # 训练路径航点
    # hangdian1_lon = 150.6758493
    # hangdian1_lat = -33.7219811
    hangdian1_lon = 150.6769266
    hangdian1_lat = -33.7221793
    hangdian2_lon = 150.6769266
    hangdian2_lat = -33.7221793
    # 设置初始点经纬度
    origin_lon = hangdian1_lon
    origin_lat = hangdian1_lat
    origin_x = gnc_coordinate.Lon_trans_to_E(hangdian1_lon, hangdian1_lat , origin_lon, origin_lat)
    origin_y = gnc_coordinate.Lat_trans_to_N(hangdian1_lon, hangdian1_lat , origin_lon, origin_lat)
    hangdian2_x = gnc_coordinate.Lon_trans_to_E(hangdian2_lon, hangdian2_lat , origin_lon, origin_lat)
    hangdian2_y = gnc_coordinate.Lat_trans_to_N(hangdian2_lon, hangdian2_lat , origin_lon, origin_lat)

    def update_func_gps(handle, lat, lon, alt):
        # 将ROS传回的经纬度转换为坐标系
        boat_x = gnc_coordinate.Lon_trans_to_E(lon, lat, origin_lon, origin_lat)
        boat_y = gnc_coordinate.Lat_trans_to_N(lon, lat, origin_lon, origin_lat)
        boat_pos_queue.put((boat_x, boat_y))


    def update_func_ea(handle, yaw, pitch, roll):
        boat_yaw_queue.put(yaw)

    def update_func_ground_speed(handle, speed , forward_v , lateral_v):
        boat_vel_queue.put((forward_v,lateral_v))
        boat_speed_queue.put(speed)

    def update_func_gyro(handle, gyroX, gyroY, gyroZ):
        pass

    def update_func_acc(handle, accX, accY, accZ):
        pass

    rospy.init_node('my_ROSVRX_node')

    update_handle = None

    ros_vrx_info = ROSVRXInfo(update_func_gps, update_func_ea, update_func_ground_speed,
                              update_func_gyro, update_func_acc, update_handle)

    propeller_mid = 0
    servo_mid = 0
    rosvrx_control = propulsion_VRX(propeller_mid, servo_mid)

    # 在主函数中调用 manual_act 控制无人船

    def publish_commands():
        import os
        if not os.path.exists("4_load_data.txt"):
            with open("data.txt", "w") as f:
                f.write("三级海况下TD3跟踪\n")

        while not rospy.is_shutdown():
            boat_coordinates = []
            dis_error =[]
            angs = []
            angs_error = []
            power_left = []
            power_right = []
            servo = []
            # 下载模型参数 并放到模型中
            state_dim = 7  # 根据实际情况设置状态维度
            action_dim = 3  # 根据实际情况设置动作维度
            hidden_dim = 128  # 根据实际情况设置隐藏层维度
            loaded_model = PolicyNetwork(state_dim, action_dim, hidden_dim)
            loaded_model.load_state_dict(torch.load('/home/ltp/train model/4/net-weight3.21.pt'))
            # 将模型设置为评估模式
            loaded_model.eval()
            # 提取船当前的状态

            goto_point([-260,227])
            time.sleep(1)
            boat_x, boat_y = boat_pos_queue.get()
            print(boat_x, boat_y)
            first_time =time.time()
            hd_list = [(0,0),(0,50),(20,50),(20,0),(40,0),(40,50),(60,50),(60,0)]
            for i, coordinates in enumerate(hd_list):
                x, y = coordinates
                if i > 0 :
                    prev_x, prev_y = hd_list[i - 1]
                else:
                    prev_x, prev_y = origin_x, origin_y
                dt_ang_error = 0
                dt_t_error = 0
                pre_ang_error = 0
                pre_t_error = 0
                action = [0, 0, 0]
                # time.sleep(10)
                pre_time = time.time()
                while (x- boat_x)**2+(y - boat_y)**2 >25:
                    boat_x, boat_y = boat_pos_queue.get()
                    boat_coordinates.append((boat_x, boat_y))
                    yaw = boat_yaw_queue.get()
                    speed = boat_speed_queue.get()
                    forward_v, lateral_v = boat_vel_queue.get()
                    t_error, ang_error = calculate_metrics(yaw, forward_v, lateral_v, prev_x, prev_y,
                                                           x, y, boat_x, boat_y)
                    # input = torch.from_numpy(np.array([t_error, ang_error, dt_t_error, dt_ang_error] + action)).double()
                    input = np.array([t_error, ang_error, dt_t_error, dt_ang_error] + action)

                    input = torch.FloatTensor(input).unsqueeze(0).to(device)
                    with torch.no_grad():
                        output_action = loaded_model(input)
                    output_action = np.squeeze(output_action.numpy())
                    power_L = (output_action[0] + 1)/2
                    power_R = (output_action[1] + 1)/2
                    turn = output_action[2]
                    rosvrx_control.manual_act(power_L, power_R, turn)
                    with open("4_load_data.txt", "a") as f:
                        # 这里的 angle_error = next_state[0], angle_error_deriv = next_state[1]
                        f.write(f"速度: {speed}, 路径误差: {t_error}, 角度误差(度): {ang_error * 180 / math.pi}, 当前位置: {boat_x},{boat_y},左推进: {power_L},右推进: {power_R},舵角: {turn}\n")
                    current_time = time.time()
                    dis_error.append((current_time - first_time,t_error))
                    power_left.append((current_time - first_time, power_L))
                    power_right.append((current_time - first_time, power_R))
                    servo.append(((current_time - first_time), turn))
                    angs.append(((current_time - first_time),yaw))
                    angs_error.append(((current_time - first_time),ang_error))
                    dt_t_error = (t_error - pre_t_error) / (current_time - pre_time)
                    dt_ang_error = (ang_error - pre_ang_error) / (current_time - pre_time)
                    pre_t_error = t_error
                    pre_ang_error = ang_error
                    pre_time = current_time
                    action = [output_action[0],output_action[1],output_action[2]]
            # 从 angs_error 中提取 ang_error 列表
            ang_errors = [ang_error for _, ang_error in angs_error]
            # 使用 NumPy 计算方差
            variance = np.var(ang_errors)
            print("方差为:", variance)
            plot_boat_coordinates(boat_coordinates,'无人船位置图')
            plot_boat_coordinates(dis_error,'无人船误差图')
            plot_boat_coordinates(angs_error,'无人船角度误差图')
            plot_boat_coordinates(power_left,'无人船图')
            plot_boat_coordinates(power_right, '无人船图')
            plot_boat_coordinates(servo,'无人船舵角输出图')
    # 创建一个线程来发布命令
    publish_thread = threading.Thread(target=publish_commands)
    publish_thread.start()

    try:
        ros_vrx_info.ROSVRX_info_create()

    except rospy.ROSInterruptException:
        pass
